package org.example.avl;

import static java.lang.Long.rotateLeft;
import static java.lang.Long.rotateRight;

public class AvlTree {
    // 使用静态内部类,避免了使用外部类创建实例
    // 使用外部类的实例才能创建内部类
    // 使用静态内部类能够直接创建实例
    static class TreeNode{
        public int val;
        public int bf;// 平衡因子
        public TreeNode left;
        public TreeNode right;
        public TreeNode parent;

        public TreeNode(int val) {
            this.val = val;
        }
    }

    public TreeNode root;// 这棵树的根节点
    // 先写一颗二叉搜索树

    // 插入节点
    public boolean insert(int val){
        TreeNode newroot = new TreeNode(val);
        // 节点为空
        if(root == null){
            root = newroot;
            return true;
        }
        // 节点不为空
        TreeNode cur =root;
        TreeNode parent = null;
        while(cur != null){
            // 在叶子节点的位置插入一个新的节点
            if(cur.val < val){
                parent = cur;
                cur = cur.right;
            }else if(cur.val == val){
                return false;
            }else{
                parent = cur;
                cur = cur.left;
            }
        }
        // cur == null
        if(val > parent.val){
            parent.right = newroot;
        }else{
            parent.left = newroot;
        }
        newroot.parent = parent;
        cur = newroot;
        // 检查平衡因子

        while(parent != null) {
            if (cur == parent.right) {
                parent.bf++;
            }else if(cur == parent.left){
                parent.bf--;
            }

            if(parent.bf == 0){
                // 已经平衡了
                break;
            }else if(parent.bf == 1 || parent.bf == -1){
                // 继续向上去修改平衡因子
                cur = parent;
                parent = cur.parent;
            }else{
                if(parent.bf == 2){
                    if(cur.bf == 1){
                        // 右边高,左单旋
                        rotateLeft(parent);
                    }
                    if(cur.bf == -1){
                        // 右左双旋
                        rotateRL(parent);
                    }
                }else if(parent.bf == -2){
                    if(cur.bf == -1){
                        // 左边高,右单旋
                        rotateRight(parent);
                    }
                    if(cur.bf == 1){
                        // 左右双旋
                        rotateLR(parent);
                    }
                }
            }
        }

        return true;
    }

    private void rotateRL(TreeNode parent) {
        TreeNode subR = parent.left;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;

        rotateRight(parent.right);
        rotateRight(parent);

        if(bf == 1){
            parent.bf = -1;
            subRL.bf = 0;
            subR.bf = 0;
        }else if(bf == -1){
            parent.bf = 0;
            subRL.bf = 0;
            subR.bf = -1;
        }
    }

    private void rotateLR(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf;

        rotateLeft(parent.left);
        rotateRight(parent);

        if(bf == -1){
            subLR.bf = 0;
            subL.bf = 0;
            parent.bf = 1;
        }else if(bf == 1){
            subLR.bf = 0;
            subL.bf = -1;
            parent.bf = 0;
        }else{
            // bf == 0
            return;
        }
    }

    private void rotateRight(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;


        parent.left = subLR;
        // subLR可能不存在
        if(subLR != null) {
            subLR.parent = parent;
        }

        TreeNode pParent = parent.parent;
        if(parent == root){
            subL = root;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subL;
            }else{
                pParent.right = subL;
            }
        }
        subL.parent = pParent;
        subL.right = parent;
        parent.parent = subL;

        subL.bf = 0;
        parent.bf = 0;
    }

    private void rotateLeft(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;

        subR.left = parent;
        TreeNode pParent = parent.parent;
        // 可能这个节点不是root,还有其他的节点
        if(parent == root){
            root = subR;
            root.parent = null;
        }else{
            if(pParent.left == parent){
                pParent.left = subR;
            }else{
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        parent.parent = subR;

        parent.right = subRL;
        if(subRL != null){
            subRL.parent = parent;
        }

        subR.bf = 0;
        parent.bf = 0;
    }
}
